package com.bins.langchain.easy.chroma.controller;

import com.bins.langchain.easy.chroma.service.ChromaService;
import jakarta.annotation.Resource;
import org.apache.commons.lang3.ArrayUtils;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.Arrays;
import java.util.List;

@RestController
public class EmbeddingController {

    @Resource
    private OllamaEmbeddingModel ollamaEmbeddingModel;

    @Autowired
    private ChromaService chromaService;

    @GetMapping("/text")
    public List<Float> textEmbedding(String text) {
        float[] floatArray = ollamaEmbeddingModel.embed(text);
        Float[] objectArray = ArrayUtils.toObject(floatArray);
        List<Float> floatList = Arrays.asList(objectArray);
        return floatList;
    }

    @PostMapping("/save")
    public void saveEmbedding() {
        chromaService.saveEmbedding();
    }

    @GetMapping("/search")
    public List<String> search() {
        return chromaService.search();
    }

}